package com.treasure.infrastructure.domain.extraction.vector;

import com.alibaba.dashscope.exception.NoApiKeyException;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.treasure.common.enums.TrainingStatEnum;
import com.treasure.common.exception.ParameterException;
import com.treasure.common.exception.ServiceException;
import com.treasure.common.model.UploadDocumentRequest;
import com.treasure.common.types.Corp;
import com.treasure.common.types.FileName;
import com.treasure.domain.domain.DocumentIndexStoreDomainService;
import com.treasure.domain.domain.entity.document.UploadFileEntity;
import com.treasure.infrastructure.index.IndexDocHandleProvider;
import com.treasure.infrastructure.index.IndexDocHandler;
import com.treasure.infrastructure.index.IndexDocTypeEnum;
import com.treasure.infrastructure.model.FileSegmentBean;
import com.treasure.infrastructure.model.po.KnowledgeDocument;
import com.treasure.infrastructure.model.po.KnowledgeDocumentVector;
import com.treasure.infrastructure.repository.dao.mapper.KnowledgeDocumentMapper;
import com.treasure.infrastructure.repository.dao.mapper.KnowledgeDocumentVectorMapper;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionDefinition;
import org.springframework.util.CollectionUtils;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

import static dev.langchain4j.model.openai.OpenAiModelName.GPT_3_5_TURBO;

/**
 * @author sunhongmin
 * @date 2023/10/21 10:47
 * @description
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class DocumentIndexStoreDomainServiceImpl implements DocumentIndexStoreDomainService {
    public static final int MAX_OVERLAP_SIZE_IN_TOKENS = 30;
    public static final int MAX_SEGMENT_SIZE_IN_TOKENS = 400;
    public static final String LINE_BREAK = "\n";
    private final KnowledgeDocumentMapper knowledgeDocumentMapper;
    private final IndexDocHandleProvider indexDocHandleProvider;
    private final KnowledgeDocumentVectorMapper knowledgeDocumentVectorMapper;
    private final PlatformTransactionManager platformTransactionManager;

    @Override
    public void indexDocumentAndStore(UploadDocumentRequest uploadDocumentRequest, File file, FileName fileName) {
        KnowledgeDocument knowledgeDocument = knowledgeDocumentMapper.selectById(uploadDocumentRequest.getDocumentId());
        if (Objects.isNull(knowledgeDocument)) {
            throw new ParameterException("doDocumentIndexStore error 文档不存在 , documentId=" + uploadDocumentRequest.getDocumentId());
        }
        TransactionStatus transaction = platformTransactionManager.getTransaction(new DefaultTransactionDefinition());
        try {
            doIndexTransactionMethod(uploadDocumentRequest, file, fileName, knowledgeDocument);
            platformTransactionManager.commit(transaction);
        } catch (Exception e) {
            log.info("indexDocumentAndStore error ", e);
            knowledgeDocument.setTrainingStat(TrainingStatEnum.FAIL_TRAINING.getValue());
            String message = e.getMessage();
            if (StringUtils.isNotEmpty(message) && message.length() > 30) {
                message = message.substring(0, 30) + "...";
            }
            knowledgeDocument.setTrainingMsg("文档训练失败，message=" + message);
            platformTransactionManager.rollback(transaction);
        }

        knowledgeDocumentMapper.updateById(knowledgeDocument);
    }

    public void doIndexTransactionMethod(UploadDocumentRequest uploadDocumentRequest, File localFile, FileName fileName, KnowledgeDocument knowledgeDocument) throws IOException, NoApiKeyException {
        Document document = FileSystemDocumentLoader.loadDocument(localFile.toPath());
        DocumentSplitter documentSplitter = DocumentSplitters
                .recursive(MAX_SEGMENT_SIZE_IN_TOKENS, MAX_OVERLAP_SIZE_IN_TOKENS, new OpenAiTokenizer(GPT_3_5_TURBO));
        List<TextSegment> segments = documentSplitter.split(document);
        if (CollectionUtils.isEmpty(segments)) {
            throw new ServiceException("切分文档失败，段落内容为空");
        }

        List<String> paragraphList = segments.stream().map(TextSegment::text).collect(Collectors.toList());
        // 自定义语意连贯处理
//        List<String> handledParagraphList = paragraphOverLoad(paragraphList); // 暂时注释

        // 删除旧vector
        deleteDocumentOldVector(new Corp(uploadDocumentRequest.getCorpId()), knowledgeDocument);

        // 存储文本向量
        IndexDocHandler indexHandler = indexDocHandleProvider.getInstance(IndexDocTypeEnum.INDEX_VECTOR_DOC);
        FileSegmentBean fileSegmentBean = new FileSegmentBean();
        fileSegmentBean.setFileName(fileName.getValue());
        fileSegmentBean.setSegmentList(paragraphList);
        indexHandler.indexResourceSegment(uploadDocumentRequest, Collections.singletonList(fileSegmentBean));
        knowledgeDocument.setTrainingStat(TrainingStatEnum.SUCCESS_TRAINING.getValue());
        knowledgeDocument.setTrainingMsg("文档训练成功");
    }


    /**
     * 语意连贯处理
     * 如果段尾换行符后的文本小于重叠限制，则添加至下一段首部
     *
     * @param paragraphList
     * @return
     */
    public static List<String> paragraphOverLoad(List<String> paragraphList) {
        if (paragraphList == null || paragraphList.isEmpty()) {
            return Collections.emptyList();
        }

        List<String> updatedList = new ArrayList<>(paragraphList.size());
        updatedList.add(paragraphList.get(0)); // 添加首段

        for (int i = 1; i < paragraphList.size(); i++) {
            String previousParagraph = paragraphList.get(i - 1);
            List<String> previousParagraphLines = Arrays.stream(previousParagraph.split(LINE_BREAK))
                    .filter(StringUtils::isNotBlank)
                    .collect(Collectors.toList());

            if (!CollectionUtils.isEmpty(previousParagraphLines) &&
                    previousParagraphLines.size() >= 2 &&
                    previousParagraphLines.get(previousParagraphLines.size() - 1).length() <= MAX_OVERLAP_SIZE_IN_TOKENS) {
                String newParagraph = previousParagraphLines.get(previousParagraphLines.size() - 1) +
                        LINE_BREAK + paragraphList.get(i);
                updatedList.add(newParagraph);
            } else {
                updatedList.add(paragraphList.get(i));
            }
        }

        return updatedList;
    }

    private void deleteDocumentOldVector(Corp corp, KnowledgeDocument documentExists) {
        log.info("checkAndGetDocument delete old document vector documentId:{}", documentExists.getId());
        QueryWrapper<KnowledgeDocumentVector> documentVectorQuery = new QueryWrapper<>();
        documentVectorQuery.eq("corp_id", corp.getValue())
                .eq("document_id", documentExists.getId());
        List<KnowledgeDocumentVector> knowledgeDocumentVectors = knowledgeDocumentVectorMapper.selectList(documentVectorQuery);
        if (!CollectionUtils.isEmpty(knowledgeDocumentVectors)) {
            knowledgeDocumentVectors.stream().map(it -> {
                KnowledgeDocumentVector knowledgeDocumentVector = new KnowledgeDocumentVector();
                knowledgeDocumentVector.setId(it.getId());
                knowledgeDocumentVector.setHasDeleted(Boolean.TRUE);
                return knowledgeDocumentVector;
            }).forEach(knowledgeDocumentVectorMapper::updateById);
        }
    }


}

